#!/usr/bin/python

# This is the ANCIENT program. This program will find extremely old repeat units guided by BLAST scores to determine the correct
# Repeat unit, this program will then confirm the replication by constructing an ancestral sequence determined by a multiple alignment
# of Just TMS regions including their hydrophillic loops. Ancestral sequences are then crossed and GSAT is used to determine
# Their Z-Score.

import blast
import pickle
import hmmtop
import gsat
import tcdb
import os,sys
import re
import numpy
import threading
import ancescon
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from Bio.Alphabet import IUPAC
from math import sqrt
import shutil
from Bio.Align.Applications import ClustalwCommandline

class build():
	
	directory = "/Volumes/ProjectV/ANCIENT2/"
	dbname = "ancient.db"
	family = "2.A.1"
	split = 6
	tms_requirement = 12
	
	def __init__(self):
		
		# Declare constants for this run
		self.accessions = [] 
		self.flank = [5,5]
		self.tms =[]
		self.loops = {}
		
	def get_fastas(self):
		pb = blast.tools()
		for acc in self.accessions:
			print acc
			pb.blastp(acc)
		pb.build_xml()
		pb.maketable(0.7)
		self.blast_status = pb.status
		self.fasta_file = pb.fasta_file
		shutil.copyfile(self.fasta_file.name,"/Volumes/ProjectV/newfasta.faa")
		return
		
	def hmmtop(self):
		hm = hmmtop.tools()
		hm.add_library("SF", self.fasta_file.name)
		hm.scan_libraries()
		self.tms = hm.results['SF']
		return
		
	def get_groups(self,keys):
		db = {}
		for size in keys:
			for row in keys:
				end = row+(size-1)
				if end > len(keys):
					break
				db.setdefault(size,[]).append(range(row,end+1))
		return db
		
	def build_flanks(self):
		fastas = SeqIO.to_dict(SeqIO.parse(self.fasta_file,'fasta'))
		for symbol, tmss in self.tms.items():
			for tms, residues in tmss.items():
				prev = tmss[tms-1] if tmss.has_key(tms-1) else False
				next = tmss[tms+1] if tmss.has_key(tms+1) else False
				this = self.expand(prev,residues,next,len(fastas[symbol].seq))
				loop = fastas[symbol].seq[this[0]-1:this[1]]
				self.loops.setdefault(symbol,[]).append(loop)				
				
	def expand(self,tmp,tmt,tmn,length):
		if tmp is False:
			# This is the first tms
			if self.flank is False:
				# Use the whole len, no flank specified
				tmt[0] = 1
			else:
				if tmt[0] - self.flank[0] < 1:
					tmt[0] = 1
				else:
					tmt[0] -= self.flank[0]
		else:
			# Not the first TMS, getting new start
			if self.flank is False:
				# No left flank specified, find midpoint
				left = (tmp[1]+tmt[0])
				if left%2==0:
					# Its even, use this value
					tmt[0]=left/2
				else:
					# Its odd, round up
					tmt[0]=int(round(left/2))
			else:
				# Flanking zones are specified
				tmt[0] -= self.flank[0]
				
		# Done With start position, now calculate end position
		if tmn is False:
			# This is the last tms position
			if self.flank is False:
				# Use the remaining len of sequenec
				tmt[1] = length
			else:
				if tmt[1] + self.flank[1] > length:
					tmt[1] = length
				else:
					tmt[1] += self.flank[1]
		else:
			# This is not the last TMS, getting new END position
			if self.flank is False:
				# No right flank, find midpoint, round DOWN
				right = (tmt[1]+tmn[0])
				if right%2==0:
					tmt[1] = (right/2)-1
				else:
					tmt[1] = int(right/2)
			else:
				# Flanking zones are specified
				tmt[1] += self.flank[1]
		return tmt
		
	def calculate_stats(self):
		length = []
		data = [i for i in self.loops.values() if len(i) is self.tms_requirement]
		for entry in data:
			lens = [len(str(i)) for i in entry]
			length.extend(lens)
		ave= numpy.average(length)
		std= numpy.std(length)
		sem= std/sqrt(ave)
		self.valid_range = range(int(round(ave*2-sem)),int(round(ave*2+sem)))
			
	def write_fastas(self):
		for symbol,seq in self.loops.items():
			try:
				groups = self.get_groups(range(1,len(seq)+1))[self.split]
			except:
				pass
			for group in groups:
				filename = "%i|%s|%s.faa" %(len(seq),','.join([str(i) for i in group]),self.family)
				handle=open(self.directory+filename,"ab+")
				try:
					string = ''.join([str(self.loops[symbol][x-1]) for x in group])
				except:
					pass
				#if len(string) in self.valid_range:
				record = SeqRecord(Seq(string,IUPAC.protein),id=symbol, name=symbol)
				SeqIO.write(record,handle,'fasta')
	
	def align(self):
		files = os.listdir(self.directory)
		myfiles = [i for i in files if re.match(str(self.tms_requirement)+'\|[0-9,]+\|+'+self.family+'\.faa$',i)]
		for thisfile in myfiles:
			threading.Thread(target=self.subalign,args=(thisfile,)).start()
		for thread in threading.enumerate():
			if thread is not threading.currentThread():
				thread.join()
		return

	def subalign(self,infile):
		clustal=ClustalwCommandline("clustalw2")
		clustal.infile="'"+self.directory+infile+"'"
		clustal.outfile="'"+self.directory+infile+".aln'"
		clustal()
		return
		
	def build_ancestors(self):
		handles = []
		files = os.listdir(self.directory)
		myfiles = [i for i in files if re.match(str(self.tms_requirement)+'\|[0-9,]+\|+'+self.family+'\.faa.aln$',i)]
		for aln in myfiles:
			handles.append(self.subancescon(aln))
		for h in handles:
			out=h.communicate()
			
	def collect(self):
		self.seqs=[]
		files = os.listdir(self.directory)
		myfiles = [i for i in files if re.search('aln.ANCEINT',i)]
		old = re.compile('Root_MidPoint\s+([A-Z,-]+)')
		myfiles.sort(reverse=False, key=lambda x: int(x.split('-')[1]))
		for an in myfiles:
			temp = open(self.directory+an).read()
			seq = old.search(temp).groups()[0]
			tms = re.search('\d-([0-9,-]+)-2-A-1',an).groups()[0]
			self.seqs.append(seq)
		
	def compare_seqs(self):
		scores = []
		for i,tms in enumerate(self.seqs):
			try:
				subject = tms
				target = self.seqs[i+self.split]
				result = gsat.compare(subject,target)
				row = [str(i)+"&"+str(i+self.split),result.zscore]
				print row
				scores.append(row)
			except:
				pass
			
	def subancescon(self,aln):
		mygroup = re.sub('\.|\||,','-',aln)
		return ancescon.predict(self.directory+aln,mygroup+".ANCIENT",self.directory+mygroup+".ANCEINT")
		
		
	def intercompare(self):
		# Extract TMS requirements
		results = {}
		specific_loops={}
		tscore = 0
		myloops = [i for i in self.loops.items()]
		rh = open("anc_results.txt",'wb+')
		for symbol, tmss in myloops:
			keys = range(1,len(tmss)+1)
			try:
				groups = self.get_groups(keys)[self.split]
			except:
				continue
			for i,tms in enumerate(groups):
				subject = tms
				target = groups[i+self.split] if (i+self.split) <= len(groups)-1 else False
				if target is False:
					continue
				subject_seq = ''.join([str(self.loops[symbol][i-1]) for i in subject])
				target_seq = ''.join([str(self.loops[symbol][i-1]) for i in target])
				score = gsat.compare(subject_seq,target_seq)
				subject = [str(i) for i in subject]
				target = [str(i) for i in target]
				zones = '-'.join(subject)+' && '+'-'.join(target)
				line = "%s\t%s\t%s" %(symbol,zones,str(score.zscore))
				print line
				rh.write(line+"\n")
				#tscore =+ zscore if zscore > 0
				results.setdefault(zones,[]).append(score.zscore)
		#ic = open(make.directory+'ic.data','wb+')
		#pickle.dump(results, ic)
		#print results
				
		
if __name__ == '__main__':
	'''
	accessions = tcdb.define_family("2.A.3")
	for subject in accessions:
		try:
			make = build()

			make.accessions = [subject]
			make.get_fastas()
			make.hmmtop()
			make.build_flanks()
			#make.calculate_stats()
			make.write_fastas()

			make.align()
			make.build_ancestors()
			make.collect()
			#make.compare_seqs()
			
		except:
			pass
		finally:
			print "Finished"
			#os.system("rm "+make.directory+"*")
	'''


	make=build()
	#make.split=sys.argv[1]
	make.accessions=tcdb.define_family("2.A.53")
	make.get_fastas()
	#myfile = sys.argv[2]
	#make.fasta_file = open('/Users/saierlab/Desktop/Biotools/Shirly/MPE95.faa','r')
	make.hmmtop()
	make.build_flanks()
	make.intercompare()
